//
// Created by hjie on 23-5-30.
//
#include "stun_message.h"
#include "./../../util/amf.h"
#include "./../../util/byte_process.h"
#include "message_attr.h"
#include <cstdint>
#include <cstdio>
#include <memory>
#include <netinet/in.h>

StunMessage::StunMessage()
{
    m_msg_type = EmBindRequest;
    m_remote_port  = 0;
    m_remote_addr  = 0;
    m_magic_cookie = 0;
}

StunMessage::~StunMessage()
{
    m_attr_map.clear();
}

void StunMessage::SetMsgType(EmStunMsgType type)
{
    m_msg_type = type;
}

ParseStatus StunMessage::ParseStunMessage(UdpBufferPacket *packet, int & use_length)
{
    if (!ParseStunHeader(packet->buffer->Data(), packet->buffer->Length()))
    {
        return EmDataNotEnoughStatus;
    }
    if (m_msg_type <= EmStunMsgUnKnown || m_msg_type > EmAllocateErrorResponse)
    {
        printf("msg type value : %d.\n", m_msg_type);
        return EmParseFailStatus;
    }
    use_length += 20;
    auto result = ParseAttrData(packet, use_length);
    return result;
}


ParseStatus StunMessage::ParseStunMessage(const uint8_t * data, int data_length, int & use_length)
{
    if (!ParseStunHeader(data, data_length))
    {
        return EmDataNotEnoughStatus;
    }
    use_length += 20;
    auto result = ParseAttrData(data, data_length, use_length);
    return result;
}

bool StunMessage::ParseBindRequestData(UdpBufferPacket *packet, int &use_length)
{
    if (m_data_length < (packet->buffer->Length() - use_length))
    {
        return false;
    }
    if (m_data_length > 0)
    {
        int length = m_data_length;
        do
        {
            int use = 0;
            auto attr = MessageAttr::ProcessMessageAttr(packet->buffer->Data() + use_length, length, use);
            if (attr == nullptr)
            {
                break;
            }
            else
            {
                m_attr_map[attr->MsgAttrType()] = std::shared_ptr<MessageAttr>(attr);
                use_length += use;
                length      = length - use;
            }
        } while (length > 0);
    }
    auto address = new AttrXorMappedAddress();
    address->ParseAddrInfo(IPPROTO_IP, packet->sai, packet->sai_len);

    m_attr_map[address->MsgAttrType()] = std::shared_ptr<MessageAttr>(address);

    printf("cookie data : %u.\n", m_magic_cookie);
    printf("transaction id : %s.\n", m_transaction_id.c_str());
    printf("decoder port : %d.\n", address->LocalPort());
    printf("decoder addr : %s.\n", address->LocalIp().c_str());
    return true;
}

ParseStatus StunMessage::ParseAttrData(UdpBufferPacket *packet, int &use_length)
{
    if (m_data_length > (packet->buffer->Length() - use_length))
    {
        // printf("transaction_id:%s and m_data_length:%d(%d) command:%d.\n",m_transaction_id.c_str(), m_data_length,((packet->buffer->Length() - use_length)), m_msg_type);
        return EmDataNotEnoughStatus;
    }
    if (m_data_length > 0)
    {
        int length = m_data_length;
        do
        {
            int use = 0;
            auto attr = MessageAttr::ProcessMessageAttr(packet->buffer->Data() + use_length, length, use);
            if (attr == nullptr)
            {
                // 要处理没有解析的attr长度-避免没有处理的attr在前面然后导致该解析的attr没有解析到, -1表示解析错误
                if (use == -1)
                {
                    printf("parse stun packet attr information error.\n");
                    use_length = length;
                    return EmParseFailStatus;
                }
                use_length += use;
                length      = length - use;
                break;
            }
            else
            {
                m_attr_map[attr->MsgAttrType()] = std::shared_ptr<MessageAttr>(attr);
                use_length += use;
                length      = length - use;
            }
        } while (length > 0);
    
        auto iter = m_attr_map.find(EmFingerPrint);
        if (iter != m_attr_map.end())
        {
            AttrCrc32 * crc32 = (AttrCrc32*)iter->second.get();
            bool result = crc32->CheckCrc32Data(packet->buffer->Data(), 20 + m_data_length - 8);
            if (!result)
            {
                // crc32 check fail
                printf("check fingerprint failed.\n");
                m_attr_map.clear();
                return EmCheckFingerPrintFail;
            }
        }
    }
    // 对于每个包都添加
    if (m_attr_map.find(EmXorMappedAddress) == m_attr_map.end())
    {
        auto address = new AttrXorMappedAddress();
        address->ParseAddrInfo(IPPROTO_IP, packet->sai, packet->sai_len);
        m_attr_map[address->MsgAttrType()] = std::shared_ptr<MessageAttr>(address);
    }
    return EmParseFinishStatus;
}

ParseStatus StunMessage::ParseAttrData(
        const uint8_t *data, int length, int &use_length, struct sockaddr *addr, socklen_t *addr_length)
{
    if (m_data_length > (length - use_length))
    {
        return EmDataNotEnoughStatus;
    }
    if (m_data_length > 0)
    {
        int attr_length = m_data_length;
        do 
        {
            int use = 0;
            auto attr = MessageAttr::ProcessMessageAttr(data + use_length, attr_length, use);
            if (attr == nullptr)
            {
                if (use == -1)
                {
                    printf("parse stun packet attr information error.\n");
                    use_length = length;
                    return EmParseFailStatus;
                }
                use_length += use;
                attr_length = attr_length - use;
            }
            else 
            {
                m_attr_map[attr->MsgAttrType()] = std::shared_ptr<MessageAttr>(attr);
                use_length += use;
                attr_length = attr_length - use;
            }

        }while(attr_length > 0);

        auto iter = m_attr_map.find(EmFingerPrint);
        if (iter != m_attr_map.end())
        {
            AttrCrc32 * crc32 = (AttrCrc32*)iter->second.get();
            bool result = crc32->CheckCrc32Data(data, length - 8);
            if (!result)
            {
                // crc32 check fail
                printf("check fingerprint failed.\n");
                m_attr_map.clear();
                return EmCheckFingerPrintFail;
            }
        }
    }
    if (addr && addr_length)
    {
        if (m_attr_map.find(EmXorMappedAddress) == m_attr_map.end())
        {
            auto address = new AttrXorMappedAddress();
            address->ParseAddrInfo(IPPROTO_IP, addr, *addr_length);
            m_attr_map[address->MsgAttrType()] = std::shared_ptr<MessageAttr>(address);
        }
    }
    return EmParseFinishStatus;
}

void StunMessage::ComposeStunPacket(WriteBuffer &data)
{
    // stun message type
    data.Write(HToBe16((uint8_t*)&m_msg_type, 2));
    // stun message length
    uint16_t message_length = 0;
    for (auto & attr : m_attr_map)
    {
        message_length += attr.second->GetAttrLength();
    }
    // fingerprint 8 byte length
    message_length += 8;

    data.Write(HToBe16((uint8_t*)&message_length, 2));
    data.Write(HToBe32((uint8_t*)&magic_cookie, 4));
    data.Write((const uint8_t*)m_transaction_id.c_str(), m_transaction_id.size());

    for (auto & attr : m_attr_map)
    {
        attr.second->ComposeAttr(data);
    }
    // add fingerprint crc32
    std::shared_ptr<AttrCrc32> crc32 = std::make_shared<AttrCrc32>();
    // crc32->ComputeCrc32(data.Data(), data.Length());
    crc32->ComposeAttr(data);
}

std::shared_ptr<MessageAttr> StunMessage::GetAttrById(MessageAttrType type)
{
    for (auto & data : m_attr_map)
    {
        if (data.first == type)
        {
            return data.second;
        }
    }
    return nullptr;
}

void StunMessage::AddStunAttr(MessageAttr *attr)
{
    if (!attr)
    {
        return;
    }
    std::shared_ptr<MessageAttr> msg_attr = std::shared_ptr<MessageAttr>(attr);
    m_attr_map[attr->MsgAttrType()] = msg_attr;
}

void StunMessage::AddStunAttr(std::shared_ptr<MessageAttr> attr)
{
    if (attr)
    {
        m_attr_map[attr->MsgAttrType()] = attr;
    }
}

bool StunMessage::ParseStunHeader(const uint8_t * data, int length)
{
    if (length < 2)
    {
        return false;
    }
    uint16_t stun_msg_type;
    memcpy((uint8_t*)&stun_msg_type, data, 2);
    if (length < 4)
    {
        return false;
    }
    uint16_t data_length = 0;
    memcpy((uint8_t*)&data_length, data + 2, 2);
    if (length < 8)
    {
        return false;
    }
    uint32_t magic_cookie;
    memcpy((uint8_t*)&magic_cookie, data + 4, 4);
    if (length < 20)
    {
        return false;
    }
    std::string transaction_id;
    transaction_id.append((char*)(data + 8), 12);

    m_msg_type       = (EmStunMsgType)be16toh(stun_msg_type);
    m_data_length    = be16toh(data_length);
    m_magic_cookie   = magic_cookie;
    m_transaction_id = transaction_id;
    return true;
}

std::string StunMessage::ProducerTransactionId()
{
    uint8_t data[6] = {0};
    char buff[128]  = {0};
    int length = 0;
    for (int index = 0; index < 6; index++)
    {
        data[index] = rand()%256;
        int use = snprintf(buff + length, sizeof(buff) - length, "%.2x", data[index]);
        length += use;
    }
    return buff;
}

uint16_t StunMessage::GetAttrLength()
{
    return m_data_length;
}

